# some of the following codes are learned from Biopython
def ProteinIterator(handler):
    # chomp the comment lines if exists
    while True:
        line = handler.readline()
        if line == '': return
        if len(line.strip().split('\t')) == 6: break

    while True:
        tmp = line.strip().split('\t')
        if len(tmp) != 6:
            raise ValueError('Records in protein file need 6 columns')
        line = handler.readline()
        yield tmp

        if not line: return

def OrthIterator(handler):
    while True:
        line = handler.readline()
        if line == '': return
        if len(line.strip().split('\t')) == 2: break

    while True:
        tmp = line.strip().split('\t')
        if len(tmp) != 2:
            raise ValueError('Records in ortholog pair file need 2 columns')
        line = handler.readline()
        yield tmp

        if not line: return

# new Iteratros can be added here
# remember add it to _FormatToIterator

_FormatToIterator = {'protein': ProteinIterator, 'orth': OrthIterator}

def parse(handler, format):

    if isinstance(handler, basestring) :
        raise TypeError("Need a file handle, not a string (i.e. not a filename)")
    if not isinstance(format, basestring) :
        raise TypeError("Need a string for the file format (lower case)")
    if not format :
        raise ValueError("Format required (lower case string)")
    if format != format.lower() :
        raise ValueError("Format string '%s' should be lower case" % format)

    if format in _FormatToIterator :
        iterator_generator = _FormatToIterator[format]
        return iterator_generator(handler)
    else:
        raise ValueError("Unknown format '%s'" % format)


def _test():
    import os
    os.chdir('..')
    os.chdir('Test')
    handler = open('protein_sample', 'r')
    for record in parse(handler, 'protein'):
        print record


if __name__ == '__main__':
    _test()




#END